﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace NMF.Optimizations.Linq
{
    internal class SelectSource<T1, T2> : ObservableSourceBase<T2>, IObservableSource<T2>, ISourceObserver<T1>
    {
        public SelectSource(IObservableSource<T1> items, Func<T1, T2> selector)
        {
            this.Source = items;
            this.Selector = selector;

            items.Subscribe(this);
        }

        private Dictionary<T1, T2> selectedDict = new Dictionary<T1, T2>();

        public T2 Select(T1 item)
        {
            T2 value;
            if (!selectedDict.TryGetValue(item, out value))
            {
                value = Selector(item);
                selectedDict.Add(item, value);
            }
            return value;
        }

        public T2 SelectAndRemove(T1 item)
        {
            T2 value;
            if (selectedDict.TryGetValue(item, out value))
            {
                selectedDict.Remove(item);
            }
            else
            {
                value = Selector(item);
            }
            return value;
        }

        void ISourceObserver<T1>.OnDeletedItem(T1 obj)
        {
            OnDeletedItem(SelectAndRemove(obj));
        }

        void ISourceObserver<T1>.OnUpdatedItem(T1 obj)
        {
            var oldSelection = Select(obj);
            var newSelection = Selector(obj);

            if (object.Equals(oldSelection, newSelection))
            {
                OnUpdatedItem(oldSelection);
            }
            else
            {
                OnDeletedItem(oldSelection);
                selectedDict[obj] = newSelection;
                OnNewItem(newSelection);
            }
        }

        void ISourceObserver<T1>.OnNewItem(T1 obj)
        {
            OnNewItem(Select(obj));
        }

        void ISourceObserver<T1>.OnReset()
        {
            OnReset();
        }

        protected override void OnReset()
        {
            selectedDict.Clear();
            base.OnReset();
        }

        public IObservableSource<T1> Source { get; private set; }
        public Func<T1, T2> Selector { get; private set; }

        public override IEnumerator<T2> GetEnumerator()
        {
            return ((IEnumerable<T1>)Source).Select(Select).GetEnumerator();
        }
    }
}
